#pragma once
#include <soft-link/thrust/system/mc/memory_resource.h>

THRUST_NAMESPACE_BEGIN

namespace system
{
namespace mc
{

namespace detail
{
    template<allocation_fn Alloc, deallocation_fn Dealloc, typename Pointer>
    class cuda_memory_resource final : public mr::memory_resource<Pointer>
    {
    public:
        Pointer do_allocate(std::size_t bytes, std::size_t alignment = THRUST_MR_DEFAULT_ALIGNMENT) override
        {
            (void)alignment;

            void * ret;
            mcError_t status = Alloc(&ret, bytes);

            if (status != mcSuccess)
            {
                mcGetLastError(); // Clear the MACA global error state.
                throw thrust::system::detail::bad_alloc(thrust::mc_category().message(status).c_str());
            }

            return Pointer(ret);
        }

        void do_deallocate(Pointer p, std::size_t bytes, std::size_t alignment) override
        {
            (void)bytes;
            (void)alignment;

            mcError_t status = Dealloc(thrust::detail::pointer_traits<Pointer>::get(p));

            if (status != mcSuccess)
            {
                thrust::mc_cub::throw_on_error(status, "MACA free failed");
            }
        }
    };

    __forceinline__ mcError_t wcudaMallocManaged(void **ptr, size_t bytes)
    {
            return ::mcMallocManaged(ptr, bytes, mcMemAttachGlobal);
    }

    __forceinline__ mcError_t wcudaMallocHost(void **ptr, size_t bytes)
    {
            return ::mcMallocHost(ptr, bytes, mcMallocHostDefault);
    }
}

}
}
THRUST_NAMESPACE_END